module Simpson

# The code is based on SciPy v1.7.1:
# https://github.com/scipy/scipy/blob/v1.7.1/scipy/integrate/_quadrature.py

export simpson

function basic_simpson(y, x=nothing, start=1, stop=size(y)[dim]-2, dx=1.0, dim=1)
    nd = length(size(y))
    slice0 = start:2:stop
    slice1 = start+1:2:stop+1
    slice2 = start+2:2:stop+2
    if x == nothing  # Even-spaced Simpson's rule.
        result = sum(collect(y[slice0] .+ 4 .* y[slice1] .+ y[slice2]), dims=dim)
        result *= dx / 3.0
    else
        # Account for possibly different spacings.
        # Simpson's rule changes a bit.
        h = diff(x, dims=dim)
        sl0 = start:2:stop
        sl1 = start+1:2:stop+1
        h0 = h[sl0]
        h1 = h[sl1]
        hsum = h0 .+ h1
        hprod = h0 .* h1
        h0divh1 = h0 ./ h1
        tmp = hsum ./ 6.0 .* (y[slice0] .* (2 .- 1.0 ./ h0divh1) .+
                              y[slice1] .* (hsum .* hsum ./ hprod) .+
                              y[slice2] .* (2 .- h0divh1))
        result = sum(tmp, dims=dim)
    end
    return convert(Float64, result[1])
end

"""
    simpson(y, x, dx, dim, even)

Integrate y(x) using samples along the given axis and the composite
Simpson's rule. If x is None, spacing of dx is assumed.
If there are an even number of samples, N, then there are an odd
number of intervals (N-1), but Simpson's rule requires an even number
of intervals. The parameter 'even' controls how this is handled.

Parameters
----------
y : array
    Array to be integrated.
x : array, optional
    If given, the points at which `y` is sampled.
dx : float, optional
    Spacing of integration points along axis of `x`. Only used when `x` is None. Default is 1.
dim : int, optional
    Axis along which to integrate. Default is the last axis.
even : String["avg", "first", "last"], optional
    'avg' : Average two results: 1) use the first N-2 intervals with
              a trapezoidal rule on the last interval and 2) use the last
              N-2 intervals with a trapezoidal rule on the first interval.
    'first' : Use Simpson's rule for the first N-2 intervals with
            a trapezoidal rule on the last interval.
    'last' : Use Simpson's rule for the last N-2 intervals with a
           trapezoidal rule on the first interval.

Notes
-----
For an odd number of samples that are equally spaced the result is
exact if the function is a polynomial of order 3 or less. If
the samples are not equally spaced, then the result is exact only
if the function is a polynomial of order 2 or less.

Examples
--------
>>> x = 0:9
>>> y = 0:9
>>> simpson(x, y)
40.5
>>> y = x .^ 3
>>> simpson(y, x)
1642.5
>>> simpson(y, x, even="first")
1644.5
>>> simpson(y, x, even="last")
1640.5
"""
function simpson(y, x=nothing, dx=1.0, dim=1; even="avg")
    nd = length(size(y))
    N = size(y)[dim]
    last_dx = dx
    first_dx = dx

    if x != nothing
        if length(size(x)) != length(size(y))
            error("If given, shape of x must be 1-D or the same as y.")
        end
        if size(x)[dim] != N
            error("If given, length of x along dim must be the same as y.")
        end
    end

    if N % 2 == 0
        val = 0.0
        result = 0.0
        even in ["avg", "last", "first"] || error("""Parameter "even" must be "avg", "last", or "first".""")
        # Compute using Simpson's rule on first intervals
        if even in ["avg", "first"]
            if isnothing(x) == false
                last_dx = x[end] - x[end-1]
            end
            val += 0.5 * last_dx * (y[end] + y[end-1])
            result = basic_simpson(y, x, 1, N-3, dx, dim)
        end
        # Compute using Simpson's rule on last set of intervals
        if even in ["avg", "last"]
            if isnothing(x) == false
                first_dx = x[2] - x[1]
            end
            val += 0.5 * first_dx * (y[2] + y[1])
            result += basic_simpson(y, x, 2, N-2, dx, dim)
        end
        if even == "avg"
            val /= 2.0
            result /= 2.0
        end
        result = result + val
    else
        result = basic_simpson(y, x, 1, N-2, dx, dim)
    end
    return result
end

end